# -*- coding: utf-8 -*-
import logging
import re
from dataclasses import dataclass, field
from typing import List, Union

import httpx
from openai import AzureOpenAI
from tenacity import retry, stop_after_attempt, wait_fixed, retry_if_not_exception_type, before_log, after_log

from videotrans.configure import config
from videotrans.configure._except import NO_RETRY_EXCEPT
from videotrans.configure.config import tr, logs
from videotrans.translator._base import BaseTrans
from videotrans.util import tools
from openai import LengthFinishReasonError

RETRY_NUMS = 3
RETRY_DELAY = 5


@dataclass
class AzureGPT(BaseTrans):
    prompt: str = field(init=False)

    def __post_init__(self):
        super().__post_init__()
        self.trans_thread = int(config.settings.get('aitrans_thread', 50))
        self.model_name = config.params.get("azure_model",'')
        self.prompt = tools.get_prompt(ainame='azure',aisendsrt=self.aisendsrt).replace('{lang}', self.target_language_name)

    @retry(retry=retry_if_not_exception_type(NO_RETRY_EXCEPT), stop=(stop_after_attempt(RETRY_NUMS)),
           wait=wait_fixed(RETRY_DELAY), before=before_log(config.logger, logging.INFO),
           after=after_log(config.logger, logging.INFO))
    def _item_task(self, data: Union[List[str], str]) -> str:
        if self._exit(): return
        model = AzureOpenAI(
            api_key=config.params.get("azure_key",''),
            api_version=config.params.get('azure_version',''),
            azure_endpoint=config.params.get("azure_api",''),
            http_client=httpx.Client(proxy=self.proxy_str)
        )
        text = "\n".join([i.strip() for i in data]) if isinstance(data, list) else data
        message = [
            {'role': 'system',
             'content': tr("You are a top-notch subtitle translation engine.")},
            {'role': 'user',
             'content': self.prompt.replace('<INPUT></INPUT>', f'<INPUT>{text}</INPUT>')},
        ]

        logs(f"\n[AzureGPT]请求数据:{message=}")
        response = model.chat.completions.create(
            model=config.params.get("azure_model",''),
            messages=message
        )
        logs(f'[AzureGPT]返回响应:{response=}')
        if not hasattr(response,'choices'):
            raise RuntimeError(str(response))
        
        if response.choices[0].finish_reason=='length':
            raise LengthFinishReasonError(completion=response)
        if response.choices[0].message.content:
            result = response.choices[0].message.content.strip()
        else:
            logs(f'[AzureGPT]请求失败:{response=}',level='warn')
            raise RuntimeError(f"[Azure] {response.choices[0].finish_reason}: {response=}")

        match = re.search(r'<TRANSLATE_TEXT>(.*?)</TRANSLATE_TEXT>', result, re.S)
        if match:
            return match.group(1)
        return result.strip()
